import cv2
import numpy as np
import os
from matplotlib import pyplot as plt

"""
异常判断：可以在根据不同阶段的图像特征，在不同阶段去做异常判断，并且可以以不同阶段的投票结果进行综合判断
当前问题：在二值化后，就会有边缘被多识别为凸起的背景，导致边缘不平整，所以在二值化前做分析？
"""

def show(img,img_name):
    cv2.imshow(img_name,img)
    cv2.waitKey(0)
    #cv2.destroyAllWindows()

def filter_edges_by_connectivity(edges, connectivity=8,min_line_length=20):
    ret,labels = cv2.connectedComponents(edges, connectivity=connectivity)  # connectivity表示连通域的连通类型，可选4或8，默认为8，上下左右和对角线八个方向
    unique_labels = np.unique(labels)
    component_sizes = [np.sum(labels == label) for label in unique_labels]

    min_size = min_line_length
    filtered_labels = [label for label, size in zip(unique_labels, component_sizes) if size >= min_size]
    #print(filtered_labels)

    conn = np.isin(labels, filtered_labels)
    edges_tobe_filtered = np.where(conn, 255, 0).astype(np.uint8)
    edges = cv2.bitwise_and(edges, edges_tobe_filtered)
    return edges


def get_four_rec_area_slide_windows(img,box,margin=10,meet=None):

    x,y,w,h = box

    # 由于pad的弧度，裁剪掉弧度部分，向上取整
    cut = int(np.ceil(margin/2))+3
    # 上
    img_up = img[max(y-3,0):y+margin,x+cut:x+w-cut]
    # 下
    img_down = img[y+h-margin:y+h+3,x+cut:x+w-cut]
    # 左
    img_left = img[y+cut:y+h-cut,max(x-3,0):x+margin]
    # 右
    img_right = img[y+cut:y+h-cut,x+w-margin:x+w+3]


    # 腐蚀
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS,(3,1))
    img_up = cv2.erode(img_up,kernel,iterations=1)
    img_down = cv2.erode(img_down,kernel,iterations=1)
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS,(1,3))
    img_left = cv2.erode(img_left,kernel,iterations=1)
    img_right = cv2.erode(img_right,kernel,iterations=1)


    # 统计每一列连续为0的个数
    img_up_col_zero = [0]*img_up.shape[1]
    img_down_col_zero = [0]*img_down.shape[1]
    # 统计每一行连续为0的个数
    img_left_row_zero = [0]*img_left.shape[0]
    img_right_row_zero = [0]*img_right.shape[0]

    # 中断列表
    interrupt_list = []
    for i in range(img_up.shape[0]):
        # 如果整行中255像素的个数超过一定比例，break
        if np.sum(img_up[i,:] == 255) > img_up.shape[1]*0.98:
            break
        for j in range(img_up.shape[1]):
            if j in interrupt_list:
                continue
            if img_up[i,j] == 0:
                img_up_col_zero[j] += 1
            else:
                interrupt_list.append(j)
    
    # 中断列表
    interrupt_list = []
    for i in range(img_down.shape[0]-1,-1,-1):
        # 如果整行中255像素的个数超过一定比例，break
        if np.sum(img_down[i,:] == 255) > img_down.shape[1]*0.98:
            break
        for j in range(img_down.shape[1]):
            if j in interrupt_list:
                continue
            if img_down[i,j] == 0:
                img_down_col_zero[j] += 1
            else:
                interrupt_list.append(j)

    # 中断列表
    interrupt_list = []
    for i in range(img_left.shape[1]):
        # 如果整列中255像素的个数超过一定比例，break
        if np.sum(img_left[:,i] == 255) > img_left.shape[0]*0.98:
            break
        for j in range(img_left.shape[0]):
            if j in interrupt_list:
                continue
            if img_left[j,i] == 0:
                img_left_row_zero[j] += 1
            else:
                interrupt_list.append(j)
    
    # 中断列表
    interrupt_list = []
    for i in range(img_right.shape[1]-1,-1,-1):
        # 如果整列中255像素的个数超过一定比例，break
        if np.sum(img_right[:,i] == 255) > img_right.shape[0]*0.98:
            break
        for j in range(img_right.shape[0]):
            if j in interrupt_list:
                continue
            if img_right[j,i] == 0:
                img_right_row_zero[j] += 1
            else:
                interrupt_list.append(j)
    
    # 统计四个列表的方差
    img_up_col_zero_var = np.var(img_up_col_zero)
    img_down_col_zero_var = np.var(img_down_col_zero)
    img_left_row_zero_var = np.var(img_left_row_zero)
    img_right_row_zero_var = np.var(img_right_row_zero)
    #print("all:",img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var)

    # 用相对较大值筛出绝对的异常
    if max(img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var) > 6:#8.7:
        return False

    # 对每个列表使用滑动窗口计算方差,用列表保存
    window = img_up.shape[1]//5
    step = window//3 if window//3 > 0 else 1
    img_up_col_zero_var_list = []
    for i in range(0,img_up.shape[1]-window,step):
        img_up_col_zero_var_list.append(np.var(img_up_col_zero[i:i+window]))
    
    img_down_col_zero_var_list = []
    for i in range(0,img_down.shape[1]-window,step):
        img_down_col_zero_var_list.append(np.var(img_down_col_zero[i:i+window]))
    
    window = img_left.shape[0]//5
    step = window//3 if window//3 > 0 else 1
    img_left_row_zero_var_list = []
    for i in range(0,img_left.shape[0]-window,step):
        img_left_row_zero_var_list.append(np.var(img_left_row_zero[i:i+window]))

    img_right_row_zero_var_list = []
    for i in range(0,img_right.shape[0]-window,step):
        img_right_row_zero_var_list.append(np.var(img_right_row_zero[i:i+window]))
    
    # 取最大的方差
    img_up_col_zero_var = max(img_up_col_zero_var_list)
    img_down_col_zero_var = max(img_down_col_zero_var_list)
    img_left_row_zero_var = max(img_left_row_zero_var_list)
    img_right_row_zero_var = max(img_right_row_zero_var_list)
    #print("seg:",img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var)


    if max(img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var) > 8.5: #13.1: #8.5:
        # 特判
        if max(img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var) == img_up_col_zero_var and img_up_col_zero_var < 14.3: #15.7:
            pass
        else:
            return False
    
    # 如果每个列表中方差小于0.25的个数超过一定比例，则认为是异常
    min_threshold = 3 
    scale_rate = 0.72
    if sum(1 for i in img_up_col_zero_var_list if i <0.25) >= (scale_rate)*len(img_up_col_zero_var_list):
        if img_up_col_zero_var >= min_threshold+1.25:
            return False
    if sum(1 for i in img_down_col_zero_var_list if i <0.25) >= (scale_rate)*len(img_down_col_zero_var_list):
        if img_down_col_zero_var >= min_threshold+0.55:
            return False
    if sum(1 for i in img_left_row_zero_var_list if i <0.35) >= (scale_rate)*len(img_left_row_zero_var_list):
        if img_left_row_zero_var >= min_threshold-0.4:
            return False
    if sum(1 for i in img_right_row_zero_var_list if i <0.25) >= (scale_rate)*len(img_right_row_zero_var_list):
        if img_right_row_zero_var >= min_threshold+0.5:
            return False
        
    # 如果均值大于一定值，则认为是异常
    scale_rate =  0.35
    min_threshold = 3.6
    if img_up_col_zero_var > min_threshold+0.2 and np.mean(img_up_col_zero_var_list) > img_up_col_zero_var*(scale_rate+0.15):
        # 如果img_up_col_zero_var_list中存在一定比例的数小于0.5
        cnt = sum(1 for i in img_up_col_zero_var_list if i < 0.5)
        if cnt > 0.15*len(img_up_col_zero_var_list):
            return False
    if img_down_col_zero_var > min_threshold and np.mean(img_down_col_zero_var_list) > img_down_col_zero_var*scale_rate:
        return False
    if img_left_row_zero_var > min_threshold and np.mean(img_left_row_zero_var_list) > img_left_row_zero_var*scale_rate:
        return False
    if img_right_row_zero_var > min_threshold and np.mean(img_right_row_zero_var_list) > img_right_row_zero_var*scale_rate:
        return False
    
    # 如果存在两个以上的值大于3.5，则认为是异常
    if sum(1 for max_var in [img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var] if max_var >= 3.5) >1:
        return False
    # 如果存在两个以上的值大于2.5，且剩下的值的和，小于0.83，则认为是异常（测试，待加）
    if sum(1 for max_var in [img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var] if max_var >= 2.5) >1:
        sigma = sum([img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var])
        for var in [img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var]:
            if var >=2.5:
                sigma -= var
        if sigma <= 0.83:
            return False

    # 分梯度来
    # 如果有一个值大于4.8，另外三个都小于1.2，则认为是异常
    if max(img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var) >= 4.8 and sum(1 for max_var in [img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var] if max_var <= 1.2) == 3:
        if max(img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var) == img_up_col_zero_var and img_up_col_zero_var < 5.5:
            pass
        else:
            var_list = [img_up_col_zero_var_list,img_down_col_zero_var_list,img_left_row_zero_var_list,img_right_row_zero_var_list]
            tmp_list = [img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var]
            max_index = tmp_list.index(max(tmp_list))
            analysis_list = var_list[max_index]
            thresh_percent = 0.72 if max_index != 0 else 0.69
            if sum(1 for var in analysis_list if var <=0.25)/len(analysis_list) >= thresh_percent:
                return False
    # 如果有一个值大于3.65，另外三个都小于1，则认为是异常
    if max(img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var) >= 3.65 and sum(1 for max_var in [img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var] if max_var <= 1) == 3:
        if max(img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var) == img_up_col_zero_var and img_up_col_zero_var < 5.5:
            pass
        else:
            var_list = [img_up_col_zero_var_list,img_down_col_zero_var_list,img_left_row_zero_var_list,img_right_row_zero_var_list]
            tmp_list = [img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var]
            max_index = tmp_list.index(max(tmp_list))
            analysis_list = var_list[max_index]
            thresh_percent = 0.72 if max_index != 0 else 0.69
            if sum(1 for var in analysis_list if var <=0.25)/len(analysis_list) >= thresh_percent:
                return False
    # 如果有一个值大于3.1，另外三个都小于0.88，则认为是异常
    if max(img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var) >= 3.1 and sum(1 for max_var in [img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var] if max_var <= 0.88) == 3:
        if max(img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var) == img_up_col_zero_var and img_up_col_zero_var < 3.4:
            pass
        else:
            var_list = [img_up_col_zero_var_list,img_down_col_zero_var_list,img_left_row_zero_var_list,img_right_row_zero_var_list]
            tmp_list = [img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var]
            max_index = tmp_list.index(max(tmp_list))
            analysis_list = var_list[max_index]
            thresh_percent = 0.72 if max_index != 0 else 0.69
            if sum(1 for var in analysis_list if var <=0.25)/len(analysis_list) >= thresh_percent:
                return False
    # 如果有一个值大于2.8，另外三个都小于0.62，则认为是异常
    if max(img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var) >= 2.8 and sum(1 for max_var in [img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var] if max_var <= 0.62) == 3:
        if max(img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var) == img_up_col_zero_var and img_up_col_zero_var < 3.4:
            pass
        else:
            var_list = [img_up_col_zero_var_list,img_down_col_zero_var_list,img_left_row_zero_var_list,img_right_row_zero_var_list]
            tmp_list = [img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var]
            max_index = tmp_list.index(max(tmp_list))
            analysis_list = var_list[max_index]
            thresh_percent = 0.72 if max_index != 0 else 0.69
            if sum(1 for var in analysis_list if var <=0.25)/len(analysis_list) >= thresh_percent:
                return False
    # 如果有一个值大于1.9，另外三个都小于0.25，则认为是异常
    if max(img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var) >= 1.9 and sum(1 for max_var in [img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var] if max_var < 0.25) == 3:
        if max(img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var) == img_up_col_zero_var and img_up_col_zero_var < 3.4:
            pass
        else:
            var_list = [img_up_col_zero_var_list,img_down_col_zero_var_list,img_left_row_zero_var_list,img_right_row_zero_var_list]
            tmp_list = [img_up_col_zero_var,img_down_col_zero_var,img_left_row_zero_var,img_right_row_zero_var]
            max_index = tmp_list.index(max(tmp_list))
            analysis_list = var_list[max_index]
            thresh_percent = 0.72 if max_index != 0 else 0.69
            if sum(1 for var in analysis_list if var <=0.25)/len(analysis_list) >= thresh_percent:
                return False
        
    return True


# 判断挡墙是否损坏
def is_damaged(img,box,margin=10):

    x,y,w,h = box
    x,y= max(x-margin,0),max(y-margin,0)
    w,h = min(w+margin*2,img.shape[1]-x),min(h+margin*2,img.shape[0]-y)

    img = img[y:y+h,x:x+w]
    img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)

    dst = cv2.adaptiveThreshold(img,255,cv2.ADAPTIVE_THRESH_MEAN_C,cv2.THRESH_BINARY,13,2)
    dst = cv2.GaussianBlur(dst,(3,3),0)
    
    box = (min(box[0],margin),min(box[1],margin),box[2],box[3])
    return get_four_rec_area_slide_windows(dst,(box),margin=margin)

# 判断挡墙是否损坏
def is_damaged(img,box,margin=10):

    x,y,w,h = box
    x,y= max(x-margin,0),max(y-margin,0)
    w,h = min(w+margin*2,img.shape[1]-x),min(h+margin*2,img.shape[0]-y)

    img = img[y:y+h,x:x+w]
    img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)

    dst = cv2.adaptiveThreshold(img,255,cv2.ADAPTIVE_THRESH_MEAN_C,cv2.THRESH_BINARY,13,2)
    dst = cv2.GaussianBlur(dst,(3,3),0)
    
    box = (min(box[0],margin),min(box[1],margin),box[2],box[3])
    return get_four_rec_area_slide_windows(dst,(box),margin=margin)


def find_pad(img_path,error_info):
    # print(img_path)
    is_pass = True  # 是否ok图片
    img = cv2.imread(img_path)
    #show(img,'img')
   
    # 先灰度化，再高斯模糊，减小计算量
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    gray = cv2.GaussianBlur(gray,(3,3),0)

    # 计算灰度图的分位数，丢弃灰度值
    q1,q2 = np.percentile(gray,[20,82])
    gray = np.clip(gray,q1,q2).astype(np.uint8)

    # 二值化
    ret,thresh = cv2.threshold(gray,0,255,cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    _,thresh = cv2.threshold(gray,ret*0.78,255,cv2.THRESH_BINARY)

    # 边缘提取
    kernel = np.ones((3,3),np.uint8)
    thresh_2 = cv2.erode(thresh,kernel,iterations=1)
    edges = cv2.Canny(thresh_2,ret*0.4,ret*0.8,apertureSize=3,L2gradient=True)

    # 膨胀 腐蚀
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS,(3,3))
    edges = cv2.morphologyEx(edges,cv2.MORPH_CLOSE,kernel)

    # 连通性分析
    edges = filter_edges_by_connectivity(edges,connectivity=8,min_line_length=(img.shape[0]+img.shape[1])*0.8)

    # 闭运算 填充空洞，连接一些断开的部分
    kernel = np.ones((3,3),np.uint8)
    #cross_kernel = cv2.getStructuringElement(cv2.MORPH_CROSS,(3,3))
    #ellipse_kernel_vertical = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3))
    edges = cv2.morphologyEx(edges,cv2.MORPH_CLOSE,kernel)

    # 寻找面积轮廓
    contours,_ = cv2.findContours(edges,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
    if len(contours) == 0:
        # print("contours == 0")
        # print(error_info[0],error_info[1],error_info[2],error_info[3])
        x,y = 10-int(error_info[0]),10-int(error_info[1])
        w,h = img.shape[1]-10+int(error_info[2])-x,img.shape[0]-10+int(error_info[3])-y  
    else:
        max_cnt = max(contours,key=cv2.contourArea)
        x,y,w,h = cv2.boundingRect(max_cnt)
        # 对box多一层判断，异常时回退到模型预测结果
        if (x == 0 and y == 0) or (abs((w+h) - (img.shape[1]-10+int(error_info[2])-x+img.shape[0]-10+int(error_info[3])-y)) >= 0.15*max(img.shape[0],img.shape[1])):  
            # print('回退到模型预测结果') 
            # print(error_info[0],error_info[1],error_info[2],error_info[3])
            x,y = 10-int(error_info[0]),10-int(error_info[1])
            w,h = img.shape[1]-10+int(error_info[2])-x,img.shape[0]-10+int(error_info[3])-y    
            # print('x,y,w,h:',x,y,w,h)\
                
                
    # x,y = 10-int(error_info[0]),10-int(error_info[1])
    # w,h = img.shape[1]-10+int(error_info[2])-x,img.shape[0]-10+int(error_info[3])-y
    # 假设框已经很准了，分别取上下左右的边缘，取出四个小矩形区域，做分析
    mar = min(w,h)//16
    is_pass = is_damaged(img,(x,y,w,h),margin=mar)
    
    return is_pass
    



# if __name__ == '__main__':
#     i = 0
#     for img_path in os.listdir('input'):
#         #if not img_path.startswith('img_'):
#         #    continue
#         #if img_path != 'test1.png':
#         #    continue
#         img_path = os.path.join('input',img_path)
#         find_pad(img_path)